{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### And welcome to Week 4, Day 3 - more LangGraph.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.graph.message import add_messages\n",
    "from dotenv import load_dotenv\n",
    "from IPython.display import Image, display\n",
    "import gradio as gr\n",
    "from langgraph.graph import StateGraph\n",
    "from langgraph.graph.message import add_messages\n",
    "from langgraph.prebuilt import ToolNode, tools_condition\n",
    "import requests\n",
    "import os\n",
    "from langchain_openai import ChatOpenAI\n",
    "from typing import TypedDict\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Our favorite first step! Crew was doing this for us, by the way.\n",
    "load_dotenv(override=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First, let's go set up LangSmith!\n",
    "\n",
    "https://langsmith.com"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Next, here is a useful function in LangChain community:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.utilities import GoogleSerperAPIWrapper\n",
    "\n",
    "serper = GoogleSerperAPIWrapper()\n",
    "serper.run(\"What is the capital of France?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now here is a LangChain wrapper class for converting functions into Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import Tool\n",
    "\n",
    "tool_search =Tool(\n",
    "        name=\"search\",\n",
    "        func=serper.run,\n",
    "        description=\"Useful for when you need more information from an online search\"\n",
    "    )\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now we can try out the tool the langchain way"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tool_search.invoke(\"What is the capital of France?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### And now let's write a tool ourselves\n",
    "\n",
    "We'll pick a familiar one"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "pushover_token = os.getenv(\"PUSHOVER_TOKEN\")\n",
    "pushover_user = os.getenv(\"PUSHOVER_USER\")\n",
    "pushover_url = \"https://api.pushover.net/1/messages.json\"\n",
    "\n",
    "def push(text: str):\n",
    "    \"\"\"Send a push notification to the user\"\"\"\n",
    "    requests.post(pushover_url, data = {\"token\": pushover_token, \"user\": pushover_user, \"message\": text})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "tool_push = Tool(\n",
    "        name=\"send_push_notification\",\n",
    "        func=push,\n",
    "        description=\"useful for when you want to send a push notification\"\n",
    "    )\n",
    "\n",
    "tool_push.invoke(\"Hello, me\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Back to the Graph from yesterday\n",
    "\n",
    "One small change - using TypedDict instead of BaseModel for the State object\n",
    "\n",
    "When we implement tools, we always need to make 2 changes to the code:\n",
    "\n",
    "1. Changes to provide the tools to OpenAI in json when we make the call\n",
    "\n",
    "2. Changes to handle the results back: look for the model staying that the finish_reason==\"tool_calls\" and then retrieve the call, run the function, provide the results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bring them together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [tool_search, tool_push]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 1: Define the State object\n",
    "class State(TypedDict):\n",
    "    messages: Annotated[list, add_messages]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 2: Start the Graph Builder with this State class\n",
    "graph_builder = StateGraph(State)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is different:\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "llm_with_tools = llm.bind_tools(tools)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 3: Create a Node\n",
    "\n",
    "\n",
    "def chatbot(state: State):\n",
    "    return {\"messages\": [llm_with_tools.invoke(state[\"messages\"])]}\n",
    "\n",
    "graph_builder.add_node(\"chatbot\", chatbot)\n",
    "graph_builder.add_node(\"tools\", ToolNode(tools=tools))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 4: Create Edges\n",
    "\n",
    "\n",
    "graph_builder.add_conditional_edges( \"chatbot\", tools_condition, \"tools\")\n",
    "\n",
    "# Any time a tool is called, we return to the chatbot to decide the next step\n",
    "graph_builder.add_edge(\"tools\", \"chatbot\")\n",
    "graph_builder.add_edge(START, \"chatbot\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 5: Compile the Graph\n",
    "graph = graph_builder.compile()\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### That's it! And, let's do this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat(user_input: str, history):\n",
    "    result = graph.invoke({\"messages\": [{\"role\": \"user\", \"content\": user_input}]})\n",
    "    return result[\"messages\"][-1].content\n",
    "\n",
    "\n",
    "gr.ChatInterface(chat, type=\"messages\").launch()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## OK it's time to add Memory!\n",
    "\n",
    "### BUT WAIT!\n",
    "\n",
    "We have this whole Graph maintaining the state and appending to the state.\n",
    "\n",
    "Why isn't this handling memory?\n",
    "\n",
    "### This is a crucial point for understanding LangGraph\n",
    "\n",
    "> A super-step can be considered a single iteration over the graph nodes. Nodes that run in parallel are part of the same super-step, while nodes that run sequentially belong to separate super-steps.\n",
    "\n",
    "\n",
    "One \"Super-Step\" of the graph represents one invocation of passing messages between agents.\n",
    "\n",
    "In idomatic LangGraph, you call invoke to run your graph for each super-step; for each interaction.\n",
    "\n",
    "The reducer handles state updates automatically within one super-step, but not between them.\n",
    "\n",
    "That is what checkpointing achieves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "memory = MemorySaver()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Steps 1 and 2\n",
    "graph_builder = StateGraph(State)\n",
    "\n",
    "\n",
    "# Step 3\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "llm_with_tools = llm.bind_tools(tools)\n",
    "\n",
    "def chatbot(state: State):\n",
    "    print(state)\n",
    "    return {\"messages\": [llm_with_tools.invoke(state[\"messages\"])]}\n",
    "\n",
    "graph_builder.add_node(\"chatbot\", chatbot)\n",
    "graph_builder.add_node(\"tools\", ToolNode(tools=tools))\n",
    "\n",
    "# Step 4\n",
    "graph_builder.add_conditional_edges( \"chatbot\", tools_condition, \"tools\")\n",
    "graph_builder.add_edge(\"tools\", \"chatbot\")\n",
    "graph_builder.add_edge(START, \"chatbot\")\n",
    "\n",
    "# Step 5\n",
    "graph = graph_builder.compile(checkpointer=memory)\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"1\"}}\n",
    "\n",
    "def chat(user_input: str, history):\n",
    "    result = graph.invoke({\"messages\": [{\"role\": \"user\", \"content\": user_input}]}, config=config)\n",
    "    return result[\"messages\"][-1].content\n",
    "\n",
    "\n",
    "gr.ChatInterface(chat, type=\"messages\").launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph.get_state(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Most recent first\n",
    "\n",
    "list(graph.get_state_history(config))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LangGraph gives you tools to set the state back to a prior point in time, to branch off:\n",
    "\n",
    "```\n",
    "config = {\"configurable\": {\"thread_id\": \"1\", \"checkpoint_id\": ...}}\n",
    "graph.invoke(None, config=config)\n",
    "```\n",
    "\n",
    "And this allows you to build stable systems that can be recovered and rerun from any prior checkpoint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### And now let's store in SQL\n",
    "\n",
    "### And this is the power of LangGraph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sqlite3\n",
    "from langgraph.checkpoint.sqlite import SqliteSaver\n",
    "\n",
    "db_path = \"memory.db\"\n",
    "conn = sqlite3.connect(db_path, check_same_thread=False)\n",
    "sql_memory = SqliteSaver(conn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Steps 1 and 2\n",
    "graph_builder = StateGraph(State)\n",
    "\n",
    "\n",
    "# Step 3\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "llm_with_tools = llm.bind_tools(tools)\n",
    "\n",
    "def chatbot(state: State):\n",
    "    print(state)\n",
    "    return {\"messages\": [llm_with_tools.invoke(state[\"messages\"])]}\n",
    "\n",
    "graph_builder.add_node(\"chatbot\", chatbot)\n",
    "graph_builder.add_node(\"tools\", ToolNode(tools=tools))\n",
    "\n",
    "# Step 4\n",
    "graph_builder.add_conditional_edges( \"chatbot\", tools_condition, \"tools\")\n",
    "graph_builder.add_edge(\"tools\", \"chatbot\")\n",
    "graph_builder.add_edge(START, \"chatbot\")\n",
    "\n",
    "# Step 5\n",
    "graph = graph_builder.compile(checkpointer=sql_memory)\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"3\"}}\n",
    "\n",
    "def chat(user_input: str, history):\n",
    "    result = graph.invoke({\"messages\": [{\"role\": \"user\", \"content\": user_input}]}, config=config)\n",
    "    return result[\"messages\"][-1].content\n",
    "\n",
    "\n",
    "gr.ChatInterface(chat, type=\"messages\").launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
